# frozen_string_literal: true

module EE
  module Gitlab
    module BackgroundMigration
      # Fixes the data integrity issue related to `namespace_id` column of the `vulnerability_reads` table
      module FixNamespaceIdsOfVulnerabilityReads
        extend ActiveSupport::Concern
        extend ::Gitlab::Utils::Override

        prepended do
          operation_name :fix_namespace_ids_of_vulnerability_reads
        end

        class Project < ApplicationRecord
          self.table_name = :projects
        end

        class VulnerabilityRead < ApplicationRecord
          include EachBatch

          self.table_name = :vulnerability_reads

          belongs_to :project, class_name: 'FixNamespaceIdsOfVulnerabilityReads::Project'

          def correct_namespace_id
            project.namespace_id
          end

          def needs_update?
            namespace_id != correct_namespace_id
          end
        end

        override :perform
        def perform
          each_sub_batch do |sub_batch|
            broken_records = sub_batch.includes(:project).select(&:needs_update?)

            broken_records.group_by(&:correct_namespace_id).each do |correct_namespace_id, records|
              vulnerability_ids = records.map(&:vulnerability_id)

              VulnerabilityRead.where(vulnerability_id: vulnerability_ids)
                               .update_all(namespace_id: correct_namespace_id)
            end
          end
        end

        # Overriding this method to be able to use Rails' eager loading.
        override :base_relation
        def base_relation
          VulnerabilityRead.where(batch_column => start_id..end_id)
        end
      end
    end
  end
end
